from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import pointfly as pf
from pointcnn import PointCNN
import tensorflow as tf

class Net(PointCNN):
    def __init__(self, points, features, is_training, setting):
        point_shape = points.get_shape()
        num_features = point_shape[1:3].num_elements()
        points = tf.reshape(points, [-1, num_features])
        points = pf.dense(points, 6144, 'pca_pts', is_training)
        points = tf.reshape(points, [-1, 2048,3])

        point_shape = features.get_shape()
        num_features = point_shape[1:3].num_elements()
        features = tf.reshape(features, [-1, num_features])
        features = pf.dense(features, 2048, 'pca_fts', is_training)
        features = tf.reshape(features, [-1, 2048, 1])

        # features = pf.dense(features, 2048, 'pca_fts', is_training)
        PointCNN.__init__(self, points, features, is_training, setting)

        fc_shape = self.fc_layers[-1].get_shape()
        num_features = fc_shape[1:3].num_elements()
        fc_flat = tf.reshape(self.fc_layers[-1], [-1, num_features])
        fc_ex = pf.dense(fc_flat, 480000, 'logits', is_training)

        fc_ex = tf.reshape(fc_ex, [-1, 60000, setting.num_class])
        # self.logits = pf.dense(self.fc_layers[-1], setting.num_class, 'logits',
        #                        is_training, with_bn=False, activation=None)
        self.logits = fc_ex
        pass

# def flatten_layer(layer):
#     # 获取输入层的形状，
#     # layer_shape == [num_images, img_height, img_width, num_channels]
#     layer_shape = layer.get_shape()
#
#     # 特征数量: img_height * img_width * num_channels
#     # 可以使用TensorFlow内建操作计算.
#     num_features = layer_shape[1:4].num_elements()
#
#     # 将形状重塑为 [num_images, num_features].
#     # 注意只设定了第二个维度的尺寸为num_filters，第一个维度为-1，保证第一个维度num_images不变
#     # 展平后的层的形状为:
#     # [num_images, img_height * img_width * num_channels]
#     layer_flat = tf.reshape(layer, [-1, num_features])
#
#     return layer_flat, num_features
#
# x =  tf.placeholder(tf.float32,[None,None,32,3])
#
# a , b =  flatten_layer(x)
# print("")
